from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Literal, Dict, Any, AsyncGenerator, Optional
from openai import OpenAI
from fastapi.responses import JSONResponse
import logging
import uvicorn
import httpx
from starlette.background import BackgroundTask
from config import settings


app = FastAPI()

class FunctionElementExtractionAPI:
    def __init__(self):
        self.client = OpenAI(
            api_key=settings.deepseek_api_key,
            base_url=settings.deepseek_base_url,
        )
        
        self.system_prompt = '''<身份>
你是一名擅长中学数学的老师，你需要根据上传的数学题目，提取出题目和问题中包含的所有函数和其他可以通过Geogebra绘制的关键元素。
并提供相应的绘制它们的GeoGebra指令，并以规定的格式返回
</身份>
<需求>
**完整性要求**：
   - 必须包含所有函数相关元素（函数/点/图像等），不得遗漏
   - 每个元素必须有明确的GeoGebra定义语句
   - 隐藏辅助元素显式定义后设置隐藏（如`SetVisibleInView(aux_point, false)`）
</需求>
<输出格式>### 1. **函数**
   - **主函数**：`f(x) = e^x + x + a`  
     - 该函数表示曲线 \( y = e^{x} + x + a \)，其中 \( a \) 是一个参数（在 GeoGebra 中需创建滑块，例如：`a = Slider(-10, 10, 1)`）。
   - **辅助函数**：
     - `g(x) = 2x + 5`  
       - 该函数表示给定直线 \( y = 2x + 5 \)。
     - `h(x) = e^x + 1`  
       - 该函数是主函数 \( f(x) \) 的导数（用于辅助理解切线条件，但在题目中未显式给出，故定义为辅助函数并隐藏）。
     - `P = (0, 5)`  
       - 该点是当 \( a = 4 \) 时的切点（根据题目条件计算得出，但作为辅助元素定义并隐藏）。

### GeoGebra 指令
- **创建参数滑块**：  
  `a = Slider(-10, 10, 1)`  
  （设置参数 \( a \) 的初始值范围，例如从 -10 到 10，步长 1）
  
- **定义主函数**：  
  `f(x) = e^x + x + a`

- **定义辅助函数**：  
  - 直线函数：`g(x) = 2x + 5`  
  - 导数函数（辅助隐藏）：`h(x) = e^x + 1`  
    `SetVisibleInView(h, 1, false)`  // 在图形视图中隐藏导数函数  
    `SetLabel(h, "derivative")`     // 可选：设置标签便于识别

- **定义辅助点**：  
  `P = Point({0, 5})`               // 定义切点 (0, 5)  
  `SetVisibleInView(P, 1, false)`   // 在图形视图中隐藏该点  
  `SetLabel(P, "tangency_point")`   // 可选：设置标签便于识别

</输出格式>'''

        logging.basicConfig(
            filename="function_element_extraction_api.log",
            level=logging.INFO,
            format="%(asctime)s - %(levelname)s - %(message)s"
        )

    async def extract_elements(self, problem: str) -> Dict[str, str]:
        """
        Extract function elements from a math problem
        """
        try:
            logging.info(f"Received input: {problem}")
            response = self.client.chat.completions.create(
                model=settings.deepseek_reasoner_model,
                messages=[
                    {"role": "system", "content": self.system_prompt},
                    {"role": "user", "content": problem},
                ],
                stream=False
            )
            reasoning_content = response.choices[0].message.reasoning_content
            elements = response.choices[0].message.content
            logging.info(f"Extracted function elements: {elements}")
            return {
                "reasoning": reasoning_content,
                "elements": elements
            }
        except Exception as e:
            logging.error(f"Error in extract_elements: {str(e)}")
            raise

    async def extract_elements_stream(self, problem: str) -> AsyncGenerator[Dict[str, Any], None]:
        """
        Stream the function element extraction process
        """
        try:
            logging.info(f"Received input for streaming: {problem[:100]}...")
            logging.info(f"Using model: {settings.deepseek_reasoner_model}")
            
            response = self.client.chat.completions.create(
                model=settings.deepseek_reasoner_model,
                messages=[
                    {"role": "system", "content": self.system_prompt},
                    {"role": "user", "content": problem},
                ],
                stream=True
            )
            
            reasoning_chunks = []
            element_chunks = []
            chunk_count = 0
            
            for chunk in response:
                chunk_count += 1
                delta = chunk.choices[0].delta
                
                if delta.reasoning_content:
                    reasoning_chunks.append(delta.reasoning_content)
                    logging.info(f"Reasoning chunk {chunk_count}: {delta.reasoning_content[:50]}...")
                    yield {"type": "reasoning", "content": delta.reasoning_content}
                    
                elif delta.content:
                    element_chunks.append(delta.content)
                    logging.info(f"Element chunk {chunk_count}: {delta.content[:50]}...")
                    yield {"type": "element", "content": delta.content}
            
            # 汇总结果
            total_reasoning = ''.join(reasoning_chunks)
            total_elements = ''.join(element_chunks)
            
            logging.info(f"Stream完成 - 总chunk数: {chunk_count}")
            logging.info(f"推理内容总长度: {len(total_reasoning)}")
            logging.info(f"元素内容总长度: {len(total_elements)}")
            logging.info(f"最终函数元素内容: {total_elements}")
            
            # 检查元素内容是否有效，如果无效则使用非流式接口作为fallback
            if not total_elements or total_elements.strip() in ["###", ""] or len(total_elements.strip()) < 10:
                logging.warning("流式接口返回的元素内容无效，使用非流式接口作为fallback")
                try:
                    fallback_result = await self.extract_elements(problem)
                    fallback_elements = fallback_result.get('elements', '')
                    if fallback_elements and len(fallback_elements.strip()) > 10:
                        logging.info(f"Fallback成功，元素内容长度: {len(fallback_elements)}")
                        yield {"type": "element", "content": fallback_elements}
                    else:
                        logging.error("Fallback也失败了")
                        yield {"type": "error", "content": "函数元素提取失败：流式和非流式接口都无法获取有效元素"}
                except Exception as fallback_error:
                    logging.error(f"Fallback失败: {fallback_error}")
                    yield {"type": "error", "content": f"函数元素提取失败: {fallback_error}"}
                    
        except Exception as e:
            logging.error(f"Error in extract_elements_stream: {str(e)}")
            yield {"type": "error", "content": str(e)}


RECEIVER_URL = "http://0.0.0.0:9000/receive_chunk_stream"

class ProblemInput(BaseModel):
    problem: str

class ElementsOutput(BaseModel):
    reasoning: str
    elements: str

@app.post("/extract_function_elements", response_model=ElementsOutput)
async def extract_function_elements(input: ProblemInput):
    try:
        logging.info(f"Received input: {input.problem}")
        api = FunctionElementExtractionAPI()
        result = await api.extract_elements(input.problem)
        return JSONResponse(content=result, status_code=200)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


def push_stream_to_receiver(problem: str):
    try:
        # 直接使用OpenAI客户端
        client = OpenAI(
            api_key="sk-b6cd134847e7430495eb999a9bb03956",
            base_url="https://api.deepseek.com"
        )
        
        system_prompt = '''<身份>
你是一名擅长中学数学的老师，你需要根据上传的数学题目，提取出题目和问题中包含的所有函数和其他可以通过Geogebra绘制的关键元素。
并提供相应的绘制它们的GeoGebra指令，并以规定的格式返回
</身份>
<需求>
**完整性要求**：
   - 必须包含所有函数相关元素（函数/点/图像等），不得遗漏
   - 每个元素必须有明确的GeoGebra定义语句
   - 隐藏辅助元素显式定义后设置隐藏（如`SetVisibleInView(aux_point, false)`）
</需求>
<输出格式>### 1. **函数**
   - **主函数**：`f(x) = e^x + x + a`  
     - 该函数表示曲线 \( y = e^{x} + x + a \)，其中 \( a \) 是一个参数（在 GeoGebra 中需创建滑块，例如：`a = Slider(-10, 10, 1)`）。
   - **辅助函数**：
     - `g(x) = 2x + 5`  
       - 该函数表示给定直线 \( y = 2x + 5 \)。
     - `h(x) = e^x + 1`  
       - 该函数是主函数 \( f(x) \) 的导数（用于辅助理解切线条件，但在题目中未显式给出，故定义为辅助函数并隐藏）。
     - `P = (0, 5)`  
       - 该点是当 \( a = 4 \) 时的切点（根据题目条件计算得出，但作为辅助元素定义并隐藏）。

### GeoGebra 指令
- **创建参数滑块**：  
  `a = Slider(-10, 10, 1)`  
  （设置参数 \( a \) 的初始值范围，例如从 -10 到 10，步长 1）
  
- **定义主函数**：  
  `f(x) = e^x + x + a`

- **定义辅助函数**：  
  - 直线函数：`g(x) = 2x + 5`  
  - 导数函数（辅助隐藏）：`h(x) = e^x + 1`  
    `SetVisibleInView(h, 1, false)`  // 在图形视图中隐藏导数函数  
    `SetLabel(h, "derivative")`     // 可选：设置标签便于识别

- **定义辅助点**：  
  `P = Point({0, 5})`               // 定义切点 (0, 5)  
  `SetVisibleInView(P, 1, false)`   // 在图形视图中隐藏该点  
  `SetLabel(P, "tangency_point")`   // 可选：设置标签便于识别

</输出格式>'''
        
        response = client.chat.completions.create(
            model="deepseek-reasoner",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": problem},
            ],
            stream=True
        )
        for chunk in response:
            delta = chunk.choices[0].delta
            data = {}
            if delta.reasoning_content:
                logging.info(f"Function reasoning content: {delta.reasoning_content}")
                data = {"type": "reasoning", "content": delta.reasoning_content}
            elif delta.content:
                logging.info(f"Function extracted elements: {delta.content}")
                data = {"type": "element", "content": delta.content}

            try:
                httpx.post(RECEIVER_URL, json=data, timeout=3.0)
            except Exception as e:
                logging.error(f"Failed to send chunk to receiver: {e}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/extract_function_elements_stream")
async def extract_function_elements_proxy(input: ProblemInput):
    logging.info(f"Received problem: {input.problem}")
    task = BackgroundTask(push_stream_to_receiver, problem=input.problem)
    return JSONResponse(content={"status": "accepted", "message": "Function streaming started to receiver"}, background=task)


if __name__ == "__main__":
    uvicorn.run("function_element_extraction_api:app", host="0.0.0.0", port=8003, reload=False) 